// Package bytecode contains types required for representing bytecode
// under evaluation, both as the result of compilation and as
// yet-to-be-evaluated lazy runtime values.
package bytecode

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"text/tabwriter"

	"tvl.fyi/tvix/eval/internal/op"
	"tvl.fyi/tvix/eval/value"
)

// Chunk represents a sequence of bytecode instructions, associated
// constants and other metadata emitted by the compiler for a code
// snippet.
type Chunk struct {
	Code      []byte
	Constants []value.NixValue
}

func (c *Chunk) PushOp(data op.Op) int {
	idx := len(c.Code)
	c.Code = append(c.Code, byte(data))
	return idx
}

func (c *Chunk) PushOpWithUint(data op.Op, arg uint64) int {
	idx := len(c.Code)
	c.Code = append(c.Code, byte(data))
	c.Code = binary.AppendUvarint(c.Code, arg)
	return idx
}

func (c *Chunk) PushOpWithUintArgs(data op.Op, args ...uint64) int {
	idx := len(c.Code)
	c.Code = append(c.Code, byte(data))
	c.Code = binary.AppendUvarint(c.Code, uint64(len(args)))

	for _, arg := range args {
		c.Code = binary.AppendUvarint(c.Code, arg)
	}

	return idx
}

func (c *Chunk) PushConstant(val value.NixValue) int {
	constantIdx := len(c.Constants)
	c.Constants = append(c.Constants, val)

	return c.PushOpWithUint(op.Constant, uint64(constantIdx))
}

// ReadOp reads the operation at the given index of the chunk. If
// there is no valid operation, op.Invalid is returned.
func (c *Chunk) ReadOp(idx int) op.Op {
	return op.FromByte(c.Code[idx])
}

// ReadUint reads the argument at the given index of the chunk (usually
// the index of the operation, plus one) and returns it with the
// number of bytes read. Calling this on an index that does not hold
// an argument is undefined. Don't do that!
func (c *Chunk) ReadUint(idx int) (uint64, int) {
	return binary.Uvarint(c.Code[idx:])
}

// PushJump pushes a jump op, allocating space for a fixed-size (uint16)
// argument. The returned index can later be passed to PatchJump.
func (c *Chunk) PushJump(data op.Op) int {
	idx := len(c.Code)
	c.Code = append(c.Code, byte(data), 0, 0)
	return idx
}

// PatchJump patches the jump instruction at the given index to point
// to the *next* instruction that will be emitted. Jump offsets are
// not encoded as varints because they require modification after the
// point where the bytes for them are allocated.
func (c *Chunk) PatchJump(idx int) {
	offset := len(c.Code) - idx

	if offset > math.MaxUint16 {
		panic("tvix bug: jump offset too large!")
	}

	binary.LittleEndian.PutUint16(c.Code[idx+1:], uint16(offset))
}

// ReadJumpOffset reads the jump offset for the given jump.
func (c *Chunk) ReadJumpOffset(idx int) uint16 {
	return binary.LittleEndian.Uint16(c.Code[idx+1:])
}

func (c *Chunk) String() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "=== chunk @ %p (%d instructions) ===\n", c, len(c.Code))

	if len(c.Code) == 0 {
		return "empty chunk"
	}

	w := tabwriter.NewWriter(&buf, 0, 8, 2, ' ', 0)

	idx := 0
	for idx < len(c.Code) {
		o := c.ReadOp(idx)

		if o == op.Constant {
			idx++
			constant, size := c.ReadUint(idx)
			idx += size
			fmt.Fprintf(w, "%#x\t%v\t(%d: %v)\t\n", idx-1, o, constant, c.Constants[constant])
			continue
		}

		switch o.HasArg() {
		case op.Uvarint:
			idx++
			arg, size := c.ReadUint(idx)
			fmt.Fprintf(w, "%#x\t%v\t(%d)\t\n", idx-1, o, arg)
			idx += size
		case op.Fixed:
			idx++
			arg := c.ReadJumpOffset(idx)
			fmt.Fprintf(w, "%#x\t%v\t(%d)\t\n", idx-1, o, arg)
			idx += 2
		case op.Multiple:
			fmt.Fprintf(w, "%#x\t%v\t", idx, o)
			idx++
			argCount, size := c.ReadUint(idx)
			idx += size

			fmt.Fprint(w, "(")
			for n := range argCount {
				arg, size := c.ReadUint(idx)
				idx += size

				if n == 0 {
					fmt.Fprint(w, arg)
				} else {
					fmt.Fprintf(w, ", %v", arg)
				}
			}
			fmt.Fprint(w, ")\t\n")
		default:
			fmt.Fprintf(w, "%#x\t%v\t(none)\t\n", idx, o)
			idx++
		}
	}

	w.Flush()
	return buf.String()
}
